package com.rrd.hive.udf.brickhouse.udf.collect;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDF;

/**
 * GroupCountUDF provides a sequence number for all rows which have the
 * same value for a particular grouping.
 * This allows us to count how many rows are in a grouping and cap them
 * off after a certain point.
 * <p/>
 * <p>
 * For example, we can cap-off the number of records per ks_uid with something
 * like
 * <p/>
 * select
 * ks_uid, val, group_count(ks_uid) as rank
 * from
 * ( select ks_uid, val from table1
 * distribute by ks_uid
 * sort by ks_uid, val ) ordered_keys
 * where group_count( ks_uid ) < 100
 */
@Description(name = "group_count",
        value = " A sequence id for all rows with the same value for a specific grouping")
public class GroupCountUDF extends UDF {
    private String lastGrouping = null;
    private int lastCount = 0;

    public Integer evaluate(String grouping) {
        // First time through ...
        if (this.lastGrouping == null) {
            this.lastGrouping = grouping;
            this.lastCount = 1;
            return 0;
        }
        if (this.lastGrouping != null && this.lastGrouping.equals(grouping)) {
            int retVal = this.lastCount;
            this.lastCount++;
            return retVal;
        } else {
            this.lastCount = 1;
            this.lastGrouping = grouping;
            return 0;
        }

    }
}